mod context;
mod switch;
mod task;

use core::borrow::BorrowMut;
use core::cell::RefCell;
use alloc::collections::binary_heap::BinaryHeap;
use crate::config::{BIG_STRIDE, PAGE_SIZE};
use crate::loader::{get_app_data, get_num_app};
use crate::mm::{MemorySet, VirtAddr};
use crate::timer::set_next_trigger;
use crate::trap::TrapContext;
pub use context::TaskContext;
use lazy_static::*;

use self::switch::__switch;
use self::task::*;

pub struct TaskManager {
    inner: RefCell<TaskManagerInner>,
}

unsafe impl Sync for TaskManager {}

struct TaskManagerInner {
    tasks: BinaryHeap<TaskControlBlock>, // 使用二叉堆, 提升寻找任务的效率
    current_task: TaskControlBlock,
}

impl TaskManager {
    fn get_current_token(&self) -> usize {
        let inner = self.inner.borrow();
        inner.current_task.get_user_token()
    }
    fn get_current_trap_ctx(&self) -> &mut TrapContext {
        let inner = self.inner.borrow();
        inner.current_task.get_trap_ctx()
    }
    fn mark_current_suspended(&self) {
        let mut inner = self.inner.borrow_mut();
        inner.current_task.task_status = TaskStatus::Ready;
    }

    fn mark_current_exited(&self) {
        let mut inner = self.inner.borrow_mut();
        inner.current_task.task_status = TaskStatus::Exited;
    }

    fn find_next_task(&self) -> Option<TaskControlBlock> {
        let mut inner = self.inner.borrow_mut();
        inner.tasks.pop()
    }

    fn run_first_task(&self) {
        let next = self.find_next_task().unwrap();
        let mut inner = self.inner.borrow_mut();
        inner.current_task = next;
        let next_task_ctx_ptr2 = inner.current_task.get_task_ctx_ptr2();
        let _unused: usize = 0;
        core::mem::drop(inner);
        set_next_trigger();
        unsafe {
            __switch(
                &_unused as *const usize,
                next_task_ctx_ptr2,
            );
        }
    }

    fn run_next_task(&self) {
        if let Some(mut next) = self.find_next_task() {
            next.task_status = TaskStatus::Running;

            let mut inner = self.inner.borrow_mut();

            let current_task_ctx_ptr2 = inner.current_task.get_task_ctx_ptr2();
            let next_task_ctx_ptr2 = next.get_task_ctx_ptr2();
            
            let stride = inner.current_task.stride;
            inner.current_task.pass += stride;
            
            core::mem::swap(&mut inner.current_task, &mut next); // 将当前任务制定为 next
            

            if next.task_status != TaskStatus::Exited {
                inner.tasks.push(next);
            }

            core::mem::drop(inner); // 释放 从 RefCell 中 借出来的可变引用
            // 因为 下面的 __switch 会切换控制流, 而 inner 的生命周期为这个函数的定义域
            // 所以在控制流被切换回来之前, inner 会一直存在
            // 这样的话, 其他的控制流中就不能借用 inner 了
            // 就会产生运行时错误
            // 所以要在切换控制流之前释放 
            set_next_trigger(); // 为下一个任务指定时间片
            unsafe {
                __switch(
                    next_task_ctx_ptr2, // 注意上面使用了 swap 函数
                    current_task_ctx_ptr2,
                    
                );
            }
            panic!("Unreachable in run next task");
        } else {
            panic!("All applications completed!");
        }
    }

    pub fn current_mmap(&self, start: usize, len: usize, port: usize) -> isize {
        use crate::mm::memory_set::MapPermission;
        use crate::mm::address::VPNRange;

        if start % PAGE_SIZE != 0 || port & 0x7 == 0 || port & !0x7 != 0 || len > 1073741824 {
            return -1;
        }

        if len == 0 {
            return 0;
        }

        let start_va: VirtAddr = start.into();
        let end_va: VirtAddr = (start + len).into();

        let vpn_range = VPNRange::new(start_va.floor(), end_va.ceil());
        for vpn in vpn_range {
            if self.inner.borrow().current_task.memory_set.find_vpn(vpn) {
                return -1;
            }
        }

        let map_perm = MapPermission::from_bits((port << 1) as u8).unwrap();
        
        let mut inner = self.inner.borrow_mut();
        inner.current_task
            .memory_set
            .insert_framed_area(start_va, end_va, map_perm | MapPermission::U);


        len as isize
    }

    pub fn current_munmap(&self, start: usize, len: usize) -> isize {
        let mut inner = self.inner.borrow_mut();
        let start_va = start.into();
        let end_va = (start + len).into();
        if inner.current_task.memory_set.munmap_area(start_va, end_va) {
            return len as isize;
        } else {
            return -1;
        }
    }
}

lazy_static! {
    pub static ref TASK_MANAGER: TaskManager = {
        info!("init TASK _MANAGER");
        let num_app = get_num_app();
        info!("num_app = {}", num_app);
        let mut tasks = BinaryHeap::new();
        for app_id in 0..num_app {
            let app_data = get_app_data(app_id);
            tasks.push(TaskControlBlock::new(
                app_data,
                app_id,
            ));
        }
        tasks.shrink_to_fit();

        let init = TaskControlBlock {
            task_ctx_ptr: 0,
            task_status: TaskStatus::Exited,
            priority: 0,
            stride: 0,
            pass: 0,
            memory_set: MemorySet::new_bare(),
            trap_ctx_ppn: 0.into(),
            base_size: 0,   
        };

        TaskManager {
            inner: RefCell::new(TaskManagerInner {
                tasks,
                current_task: init,
            }),
        }
    };
}

pub fn run_first_task() {
    TASK_MANAGER.run_first_task();
}

fn run_next_task() {
    TASK_MANAGER.run_next_task();
}

fn mark_current_suspended() {
    TASK_MANAGER.mark_current_suspended();
}

fn mark_current_exited() {
    TASK_MANAGER.mark_current_exited();
}

pub fn suspend_current_and_run_next() {
    mark_current_suspended();
    run_next_task();
}

pub fn exit_current_and_run_next() {
    mark_current_exited();
    run_next_task();
}

pub fn current_user_token() -> usize {
    TASK_MANAGER.get_current_token()
}

pub fn current_trap_ctx() -> &'static mut TrapContext {
    TASK_MANAGER.get_current_trap_ctx()
}

pub fn set_priority(priority: isize) -> isize {
    if priority <= 1 {
        -1
    } else {
        let mut inner = TASK_MANAGER.inner.borrow_mut();
        let old_pass = inner.current_task.pass;
        let old_prio = inner.current_task.priority;
        inner.current_task.priority = priority;
        inner.current_task.stride = BIG_STRIDE as isize / priority;
        inner.current_task.pass = old_pass / priority * old_prio;
        priority
    }
}

pub fn mmap(start: usize, len: usize, port: usize) -> isize {
    TASK_MANAGER.current_mmap(start, len, port)
}

pub fn munmap(start: usize, len: usize) -> isize {
    TASK_MANAGER.current_munmap(start, len)
}